# 训练模块
import argparse
import utils
import models
import torch
import data.dataset as dt
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
from tqdm import tqdm
from cfg import cfg_voc2012 as cfg

DEVICE = "cuda:0"


def loss_fn(output, target, args):
    output = output.permute(0, 2, 3, 1)
    output = output.reshape(output.size(0), output.size(1), output.size(2), 3, -1)

    mask_obj = target[..., 0] > 0
    mask_noobj = target[..., 0] == 0

    o_obj = output[mask_obj]
    non_obj = output[mask_noobj]
    t_obj = target[mask_obj]
    tno_obj = target[mask_noobj]

    # iou损失
    # 是否要给置信度损失输出后加上sigmoid
    if args.AddSigmoid_conf:
        o_obj[:, 0] = torch.sigmoid(o_obj[:, 0])
    if args.AddSigmoid_xy:
        o_obj[:, 1:3] = torch.sigmoid(o_obj[:, 1:3])

    # 置信度
    if args.conf_loss == 'mse':
        iou_loss_obj = torch.mean((o_obj[:, 0] - t_obj[:, 0]) ** 2)
        iou_loss_noobj = torch.mean((non_obj[:, 0] - tno_obj[:, 0]) ** 2)
    elif args.conf_loss == 'bce':
        pass
    # 回归损失
    if args.box_loss == 'mse':
        box_loss_obj = torch.mean((o_obj[:, 1:5] - t_obj[:, 1:5]) ** 2)
    elif args.box_loss == 'bce':
        pass

    # 分类损失
    if args.cls_loss == 'mse':
        cls_loss_obj = torch.mean((o_obj[:, 5:] - t_obj[:, 5:]) ** 2)
    elif args.cls_loss_obj == 'bce':
        pass

    if torch.isnan(iou_loss_obj):
        iou_loss_obj = 0
    if torch.isnan(iou_loss_noobj):
        iou_loss_noobj = 0
    if torch.isnan(box_loss_obj):
        box_loss_obj = 0
    if torch.isnan(cls_loss_obj):
        cls_loss_obj = 0

    loss = args.obj_weight * (iou_loss_obj + box_loss_obj + cls_loss_obj) + args.noobj_weight * iou_loss_noobj
    return loss


def train(args, logger):
    # 模型
    if args.backbone_net == 'darknet53':
        net = models.YoloV3Net()
    elif args.backbone_net == 'mobilenet':
        pass
    epoch = 0

    best_testloss = 10
    if args.weights_to_load == "darknet53.conv.74":
        models.load_darknet_weights(net, r"./weights/darknet53.conv.74")

    elif args.weights_to_load == None:
        pass
    else:
        checkpoint = torch.load(args.weights_to_load)
        net.load_state_dict(checkpoint['state_dict'])
        # optimizer_arg = checkpoint['optimizer']
        epoch = checkpoint['epoch']


        locals()['best_testloss'] = checkpoint["best_testloss"]
        logger.info(f'load {args.weights_to_load} ok || last_epoch {epoch} || last_loss {checkpoint["best_testloss"]}')

    # 是否冻结主网络
    if args.freeze_backbonenet:
        models.freeze_darknet(net)

    # 设置优化器类型
    if args.optimizer == 'adam':
        opt = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()))

    # sumwrt = SummaryWriter('./tensorboard')

    # 创造数据集
    train_data = dt.VOC2012_DATASET(r"data/voc2012")
    train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, drop_last=False, num_workers=4)

    test_data = dt.VOC2012_DATASET(r"data/voc2012", isTrain=False)
    test_loader = DataLoader(test_data, batch_size=args.batch_size, shuffle=True, drop_last=False, num_workers=4)

    # 使用cuda?
    net = net.to(DEVICE)

    for epoch_ in range(epoch, args.end_epoch):

        sum_loss = 0
        sum_testloss = 0
        for label_13, label_26, label_52, img in tqdm(train_loader):
            net.train()
            label_13 = label_13.to(DEVICE)
            label_26 = label_26.to(DEVICE)
            label_52 = label_52.to(DEVICE)
            img = img.to(DEVICE)

            out_13, out_26, out_52 = net(img)
            # print(out_13.shape)
            # 推理
            loss_13 = loss_fn(out_13, label_13, args)
            loss_26 = loss_fn(out_26, label_26, args)
            loss_52 = loss_fn(out_52, label_52, args)

            # 计算损失
            all_loss = loss_13 + loss_26 + loss_52

            opt.zero_grad()
            all_loss.backward()
            opt.step()
            sum_loss += all_loss.detach().cpu().item()

        with torch.no_grad():  # 防止梯度更新
            for label_13_t, label_26_t, label_52_t, img in tqdm(test_loader):
                net.eval()

                label_13_t = label_13_t.to(DEVICE)
                label_26_t = label_26_t.to(DEVICE)
                label_52_t = label_52_t.to(DEVICE)
                img = img.to(DEVICE)

                out_13_t, out_26_t, out_52_t = net(img)
                # print(out_13.shape)
                # 推理
                loss_13_t = loss_fn(out_13_t, label_13_t, args)
                loss_26_t = loss_fn(out_26_t, label_26_t, args)
                loss_52_t = loss_fn(out_52_t, label_52_t, args)

                # 计算损失
                all_loss = loss_13_t + loss_26_t + loss_52_t
                sum_testloss += all_loss.detach().cpu().item()

            # 优化

        sum_testloss = sum_testloss / len(test_loader)
        logger.info(f"{epoch_} || train loss {sum_loss / len(train_loader)} || test loss {sum_testloss}")
        if sum_testloss <= best_testloss:
            best_testloss = sum_testloss
            state = {'epoch': epoch_ + 1,  # 保存的当前轮数
                     'state_dict': net.state_dict(),  # 训练好的参数
                     # 'optimizer': opt.state_dict(),  # 优化器参数,为了后续的resume
                     'best_testloss': best_testloss, }  # 当前最好的精度}

            torch.save(state, f'./weights/{cfg.version}_best.pth')
            logger.info("===========Save weights===========")
            logger.info(f"got best test loss :{sum_testloss},save pt ok")
            logger.info("==================================")
        if epoch_ % 10 == 0:
            state = {'epoch': epoch_ + 1,  # 保存的当前轮数
                     'state_dict': net.state_dict(),  # 训练好的参数
                     # 'optimizer': opt.state_dict(),  # 优化器参数,为了后续的resume
                     'best_testloss': best_testloss, }  # 当前最好的精度}
            torch.save(state, f'./weights/{cfg.version}_{epoch_}.pth')

            logger.info(f"save weights epoch {epoch_}|| loss {sum_testloss},save pt ok")

    # 加载参数

    # 优化器和损失


def main():
    parser = argparse.ArgumentParser(description='Train face network')
    # general
    parser.add_argument('--weights-to-load', type=str, default=None, help='where to load weights')
    parser.add_argument('--freeze-backbonenet', type=bool, default=True,
                        help='if freeze,the weights of backbone net will not update')
    parser.add_argument('--backbone-net', default='darknet53', help='specify network')
    parser.add_argument('--data', type=str, default='VOC2012', help='*.data path')
    parser.add_argument('--end-epoch', type=int, default=1000, help='training epoch size.')
    parser.add_argument('--lr', type=float, default=0.1, help='start learning rate')
    parser.add_argument('--lr-step', type=int, default=10, help='period of learning rate decay')
    parser.add_argument('--optimizer', default='sgd', help='optimizer')
    parser.add_argument('--logout', default='train_log.log', help='where to save log')
    parser.add_argument('--AddSigmoid-conf', type=bool, default=False, help='output sel')
    parser.add_argument('--AddSigmoid-xy', type=bool, default=False, help='output sel')
    parser.add_argument('--conf-loss', type=str, default='mse', help='the type of confidence loss')
    parser.add_argument('--box-loss', type=str, default='mse', help='the type of box regression loss')
    parser.add_argument('--cls-loss', type=str, default='mse', help='the type of classficaion loss')
    parser.add_argument('--obj-weight', type=str, default='mse', help='the weight that plus with loss of object')
    parser.add_argument('--noobj-weight', type=str, default='mse', help='the weight that plus with loss of no object')
    parser.add_argument('--batch-size', type=str, default='mse', help='the weight that plus with loss of no object')
    args = parser.parse_args()
    utils.parse_config(args, cfg)

    logger = utils.Logger(args.logout).logger
    logger.log(level=20, msg='\t\n')
    logger.info("TRAIN START")
    logger.info("===========Load config===========")
    arg_str = str(args)[10:-1].split(',')
    for item in arg_str:
        logger.info(item.lstrip())
    logger.info("=================================")

    train(args, logger)
    # print(args)


if __name__ == '__main__':
    main()
